{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f8178839",
   "metadata": {},
   "source": [
    "This is part 1 of a tutorial series. We recommend following them in order, starting with [Part 0: Welcome to `musica`](0.%20Welcome%20to%20MUSICA.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf5a3b29",
   "metadata": {},
   "source": [
    "# User-Defined Reactions in `musica`\n",
    "\n",
    "So far, we've created chemical mechanisms using only Arrhenius reactions. Here we'll introduce a few more, including the powerful \"user-defined\" reaction. These reaction types give you control over how rate constants are calculated.\n",
    "\n",
    "We'll also throw in an `Emission` and `FirstOrderLoss`, which actually use the `UserDefined` type under the hood."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63424343",
   "metadata": {},
   "source": [
    "## 1. Importing Libraries\n",
    "Below is a list of the required libraries for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f81fdb7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import musica\n",
    "import musica.mechanism_configuration as mc\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import qmc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "pd.set_option('display.float_format', str) # This is done to make the arrays more readable\n",
    "np.set_printoptions(suppress=True) # This is done to make the arrays more readable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b68587",
   "metadata": {},
   "source": [
    "## 2. Base Chemical System\n",
    "Let's start with the chemical system we've been using in the tutorials to this point, with three species and two reactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50d98bb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = mc.Species(name=\"A\")\n",
    "B = mc.Species(name=\"B\")\n",
    "C = mc.Species(name=\"C\")\n",
    "species = [A, B, C]\n",
    "gas = mc.Phase(name=\"gas\", species=species)\n",
    "\n",
    "r1 = mc.Arrhenius(\n",
    "    name=\"A_to_B\",\n",
    "    A=4.0e-3,\n",
    "    C=50,\n",
    "    reactants=[A],\n",
    "    products=[B],\n",
    "    gas_phase=gas\n",
    ")\n",
    "\n",
    "r2 = mc.Arrhenius(\n",
    "    name=\"B_to_C\",\n",
    "    A=4.0e-3,\n",
    "    C=50,  \n",
    "    reactants=[B],\n",
    "    products=[C],\n",
    "    gas_phase=gas\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d06d4447",
   "metadata": {},
   "source": [
    "## 3. Defining User Defined Reactions\n",
    "This code cell defines three new reactions: one User Defined, one Emission, and one First Order Loss.</br>\n",
    "It then creates a new mechanism containing both the old reactions as well as the three new reactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6781a5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "r3 = mc.UserDefined(\n",
    "    name=\"complex_rxn\",\n",
    "    scaling_factor=1.0,\n",
    "    reactants=[A, (2, B)],\n",
    "    products=[A, (2, C), B],\n",
    "    gas_phase=gas\n",
    ")\n",
    "\n",
    "r4 = mc.Emission(\n",
    "    name=\"B_emission\",\n",
    "    scaling_factor=1.0,\n",
    "    products=[B],\n",
    "    gas_phase=gas\n",
    ")\n",
    "\n",
    "r5 = mc.FirstOrderLoss(\n",
    "    name=\"C_loss\",\n",
    "    scaling_factor=1.0,\n",
    "    reactants=[C],\n",
    "    gas_phase=gas\n",
    ")\n",
    "\n",
    "mechanism = mc.Mechanism(\n",
    "    name=\"musica_micm_example\",\n",
    "    species=species,\n",
    "    phases=[gas],\n",
    "    reactions=[r1, r2, r3, r4, r5]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc39dfd3",
   "metadata": {},
   "source": [
    "Notice that we don't specify rate constants for the User-Defined or First-Order Loss reactions, and we don't specify emission rates. This is because these values can change over time, and unlike standard rate constant types (like Arrhenius reactions), `micm` doesn't know how to calculate them based on the current temperature and pressure.\n",
    "\n",
    "So, it's our responsibility to calculate these rates and rate constants, and update them in the `micm` state, as needed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da99a490",
   "metadata": {},
   "source": [
    "## 4. Creating the Solver, State, and Latin Hypercube Sampler\n",
    "\n",
    "We'll use the Latin Hypercube Sampler to create the initial conditions, just like in the last tutorial. But, this time around we'll expand the LHS to eight dimensions: the five original dimensions (temperature, pressure, and concentrations of A, B, and C) and three new dimensions for the User-Defined, Emission, and First-Order Loss reactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3947b2c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = musica.MICM(mechanism = mechanism, solver_type = musica.SolverType.rosenbrock_standard_order)\n",
    "\n",
    "num_grid_cells = 100\n",
    "state = solver.create_state(num_grid_cells)\n",
    "\n",
    "ndim = 8\n",
    "nsamples = num_grid_cells\n",
    "\n",
    "# Create a Latin Hypercube sampler in the unit hypercube\n",
    "sampler = qmc.LatinHypercube(d=ndim)\n",
    "\n",
    "# Generate samples\n",
    "sample = sampler.random(n=nsamples)\n",
    "\n",
    "# Define bounds for each dimension (temperature, pressure, concentrations of A, B, and C, User-defined reaction, Emission, and Loss)\n",
    "l_bounds = [275, 100753.3, 0, 0, 0, 0, 0, 0] # Lower bounds\n",
    "u_bounds = [325, 101753.3, 20, 10, 20, 1, 0.5, 0.5] # Upper bounds\n",
    "\n",
    "# Scale the samples to the defined bounds\n",
    "sample_scaled = qmc.scale(sample, l_bounds, u_bounds)\n",
    "\n",
    "temperatures = sample_scaled[:, 0]\n",
    "pressures = sample_scaled[:, 1]\n",
    "concentrations = {\n",
    "    \"A\": [],\n",
    "    \"B\": [],\n",
    "    \"C\": []\n",
    "}\n",
    "concentrations[\"A\"] = sample_scaled[:, 2]\n",
    "concentrations[\"B\"] = sample_scaled[:, 3]\n",
    "concentrations[\"C\"] = sample_scaled[:, 4]\n",
    "\n",
    "state.set_conditions(temperatures, pressures)\n",
    "state.set_concentrations(concentrations)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af454f9",
   "metadata": {},
   "source": [
    "## 5. Creating and Setting the User Defined Parameters\n",
    "\n",
    "The last step in setting the conditions, is providing initial values for our new reactions. To do that, we need to provide them to the state with their unique names. The syntax for the label is `PREFIX.reaction-name`. The `PREFIX` is based on the type of reaction:\n",
    "* `USER` for User Defined reactions,\n",
    "* `EMIS` for Emission reactions, and\n",
    "* `LOSS` for First Order Loss reactions.\n",
    "\n",
    "Then `reaction-name` is the name we gave to each reaction when it was created (see above).\n",
    "\n",
    "Once we have the labels, we can use them in a call to the `state.set_user_defined_rate_parameters()` function to set the values of the reaction rates and rate constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49b15078",
   "metadata": {},
   "outputs": [],
   "source": [
    "state.set_user_defined_rate_parameters({\n",
    "    \"USER.complex_rxn\": sample_scaled[:, 5],\n",
    "    \"EMIS.B_emission\": sample_scaled[:, 6],\n",
    "    \"LOSS.C_loss\": sample_scaled[:, 7]\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d746c84f",
   "metadata": {},
   "source": [
    "## 6. Solving the System\n",
    "\n",
    "We'll solve the system as in previous tutorials, for 600 s in 10 s intervals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda25f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "concentrations_solved = []\n",
    "time_step = 10\n",
    "sim_length = 600\n",
    "curr_time = 0\n",
    "\n",
    "while curr_time <= sim_length:\n",
    "    solver.solve(state, time_step)\n",
    "    concentrations_solved.append(state.get_concentrations())\n",
    "    curr_time += time_step\n",
    "    state.set_user_defined_rate_parameters({\n",
    "        \"EMIS.B_emission\": sample_scaled[:, 6] * (1 + 0.5 * np.sin(curr_time / sim_length * 2 * np.pi))\n",
    "    })"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ca1a3c",
   "metadata": {},
   "source": [
    "Note how we update the emission rate at each time step. The other two rate constants remain constant throughout the simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad115f60",
   "metadata": {},
   "source": [
    "## 7. Visualizing the Results\n",
    "\n",
    "Now, let's prepare and plot the results, just like we've done before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da483cf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "concentrations_solved_expanded = []\n",
    "time = []\n",
    "for i in range(0, sim_length + 1, time_step):\n",
    "    for j in range(0, num_grid_cells):\n",
    "        concentrations_solved_expanded.append({key: value[j] for key, value in concentrations_solved[int(i/time_step)].items()})\n",
    "        time.append(i)\n",
    "df_expanded = pd.DataFrame(concentrations_solved_expanded)\n",
    "df_expanded = df_expanded.rename(columns = {'A' : 'CONC.A.mol m-3', 'B' : 'CONC.B.mol m-3', 'C' : 'CONC.C.mol m-3'})\n",
    "df_expanded['time.s'] = time\n",
    "df_expanded['ENV.temperature.K'] = np.repeat(temperatures[0], (sim_length/time_step + 1.0) * num_grid_cells)\n",
    "df_expanded['ENV.pressure.Pa'] = np.repeat(pressures[0], (sim_length/time_step + 1.0) * num_grid_cells)\n",
    "df_expanded['ENV.air number density.mol m-3'] = np.repeat(state.get_conditions()['air_density'][0], (sim_length/time_step + 1.0) * num_grid_cells)\n",
    "df_expanded = df_expanded[['time.s', 'ENV.temperature.K', 'ENV.pressure.Pa', 'ENV.air number density.mol m-3', 'CONC.A.mol m-3', 'CONC.B.mol m-3', 'CONC.C.mol m-3']]\n",
    "display(df_expanded)\n",
    "\n",
    "sns.lineplot(data=df_expanded, x='time.s', y='CONC.A.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.A.mol m-3')\n",
    "sns.lineplot(data=df_expanded, x='time.s', y='CONC.B.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.B.mol m-3')\n",
    "sns.lineplot(data=df_expanded, x='time.s', y='CONC.C.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.C.mol m-3')\n",
    "plt.title('Average concentration with CI over time')\n",
    "plt.ylabel('Concentration (mol m-3)')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "min_y = []\n",
    "max_y = []\n",
    "for i in range(0, sim_length + 1, time_step):\n",
    "        min_y.append({key: np.min(value) for key, value in concentrations_solved[int(i/time_step)].items()})\n",
    "        max_y.append({key: np.max(value) for key, value in concentrations_solved[int(i/time_step)].items()})\n",
    "time_x = list(map(float, range(0, sim_length + 1, time_step)))\n",
    "\n",
    "plt.fill_between(time_x, [y['A'] for y in min_y], [y['A'] for y in max_y], alpha = 0.4, label='CONC.A.mol m-3')\n",
    "plt.fill_between(time_x, [y['B'] for y in min_y], [y['B'] for y in max_y], alpha = 0.4, label='CONC.B.mol m-3')\n",
    "plt.fill_between(time_x, [y['C'] for y in min_y], [y['C'] for y in max_y], alpha = 0.4, label='CONC.C.mol m-3')\n",
    "plt.title('Concentration range over time')\n",
    "plt.ylabel('Concentration (mol m-3)')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}